import matplotlib.pyplot as plt
from collections import defaultdict
import argparse
import numpy as np

# 从文件中读取并解析区间数据的函数
def read_kmer_data(file_path):
    interval_data = defaultdict(dict)
    current_interval = None

    with open(file_path, 'r') as file:
        for line in file:
            line = line.strip()
            if line.startswith("Interval:"):
                current_interval = line.split("Interval: ")[1].strip()
            elif line and current_interval:
                parts = line.split(":")
                if len(parts) == 2:
                    kmer = parts[0].strip()
                    try:
                        count = int(parts[1].strip())
                        # 仅处理kmer长度等于5的数据
                        if len(kmer) == 5:
                            interval_data[current_interval][kmer] = count
                    except ValueError:
                        pass

    return interval_data

# 计算每个kmer在不同区间中的占比
def calculate_kmer_ratios(interval_data):
    kmer_ratios = defaultdict(list)
    # 对区间按指定顺序进行排序
    intervals = sorted(interval_data.keys(), key=lambda x: parse_interval_key(x))

    for interval in intervals:
        total_count = sum(interval_data[interval].values())
        for kmer, count in interval_data[interval].items():
            ratio = count / total_count if total_count > 0 else 0
            kmer_ratios[kmer].append(ratio)

    return kmer_ratios, intervals

# 对区间进行排序的函数
def parse_interval_key(interval):
    if interval == '[0,0]':
        return -1
    elif interval == '[100,100]':
        return 101
    else:
        # 提取区间中的数值部分，例如 (10,20] -> 10
        return int(interval.split(',')[0].replace('(', '').replace('[', '').strip())

# 找出变化最大的top5和最小的top5的kmer
def find_top_kmers_by_variation(kmer_ratios, top_n=5):
    kmer_variations = {kmer: np.var(ratios) for kmer, ratios in kmer_ratios.items()}

    # 找出变化最大的top_n个kmer
    top_max_kmers = sorted(kmer_variations.items(), key=lambda x: x[1], reverse=True)[:top_n]
    # 找出变化最小的top_n个kmer
    top_min_kmers = sorted(kmer_variations.items(), key=lambda x: x[1])[:top_n]

    return [kmer for kmer, _ in top_max_kmers], [kmer for kmer, _ in top_min_kmers]

# 绘制top5变化最大和最小的kmer的占比变化曲线
def plot_kmer_ratios(kmer_ratios, intervals, top_max_kmers, top_min_kmers):
    plt.figure(figsize=(14, 8))

    # 绘制变化最大的top5 kmer
    for kmer in top_max_kmers:
        plt.plot(intervals, kmer_ratios[kmer], label=f'Max Variation: {kmer}', linestyle='-', marker='o')

    # 绘制变化最小的top5 kmer
    for kmer in top_min_kmers:
        plt.plot(intervals, kmer_ratios[kmer], label=f'Min Variation: {kmer}', linestyle='--', marker='x')

    plt.xlabel('Methylation Interval')
    plt.ylabel('K-mer Proportion')
    plt.title('Top 5 K-mers with Maximum and Minimum Variation in Proportion Across Intervals')
    plt.xticks(rotation=45)
    plt.legend()
    plt.tight_layout()
    plt.show()

# 主函数，用于整体控制逻辑
def main():
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description="选择性读取kmer数据并绘制占比变化图")
    parser.add_argument("file_path", help="输入的kmer统计文件路径")
    args = parser.parse_args()

    # 读取kmer数据
    interval_data = read_kmer_data(args.file_path)

    # 计算每个kmer在不同区间中的占比
    kmer_ratios, intervals = calculate_kmer_ratios(interval_data)

    # 找出变化最大的top5和最小的top5的kmer
    top_max_kmers, top_min_kmers = find_top_kmers_by_variation(kmer_ratios)

    # 绘制这些kmer的占比变化曲线
    plot_kmer_ratios(kmer_ratios, intervals, top_max_kmers, top_min_kmers)

if __name__ == "__main__":
    main()
